package com.my.study.jdbc.aop;

import com.my.study.utils.SpringContextUtil;
import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.AfterThrowing;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.stereotype.Component;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.annotation.Isolation;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.support.DefaultTransactionDefinition;

import java.util.HashMap;
import java.util.Map;
import java.util.Stack;

/**
 * 多数据源事务切面
 *
 * @author: yidujun
 * @create: 2021/11/15 21:00
 */
@Component
@Aspect
public class MultiDataSourceTransactionalAspect {

    private static final ThreadLocal<Stack<Map<PlatformTransactionManager, TransactionStatus>>> THREAD_LOCAL = new ThreadLocal<>();

    @Pointcut("@annotation(com.my.study.jdbc.aop.MultiDataSourceTransactional)")
    public void point() {}

    @Before(value = "point() && @annotation(transactional)")
    public void before(MultiDataSourceTransactional transactional) {
        // 事务组
        String[] transactionManagers = transactional.transactionManagers();
        // 传播行为
        Propagation propagation = transactional.propagation();
        // 隔离级别
        Isolation isolation = transactional.isolation();

        Stack<Map<PlatformTransactionManager, TransactionStatus>> pairStack = new Stack<>();
        for (String transactionManager : transactionManagers) {
            PlatformTransactionManager manager = SpringContextUtil.getBean(transactionManager, PlatformTransactionManager.class);
            final DefaultTransactionDefinition definition = new DefaultTransactionDefinition();
            definition.setPropagationBehavior(propagation.value());
            definition.setIsolationLevel(isolation.value());
            Map<PlatformTransactionManager, TransactionStatus> transactionMap = new HashMap<>(2);
            transactionMap.put(manager, manager.getTransaction(definition));
            pairStack.push(transactionMap);
        }

        THREAD_LOCAL.set(pairStack);
    }

    @AfterReturning("point()")
    public void afterReturning() {
        // 栈顶弹出（后进先出）
        Stack<Map<PlatformTransactionManager, TransactionStatus>> pairStack = THREAD_LOCAL.get();
        while (!pairStack.empty()) {
            Map<PlatformTransactionManager, TransactionStatus> pair = pairStack.pop();
            pair.forEach(PlatformTransactionManager::commit);
        }
        THREAD_LOCAL.remove();
    }

    @AfterThrowing("point()")
    public void afterThrowing() {
        // 栈顶弹出（后进先出）
        Stack<Map<PlatformTransactionManager, TransactionStatus>> pairStack = THREAD_LOCAL.get();
        while (!pairStack.empty()) {
            Map<PlatformTransactionManager, TransactionStatus> pair = pairStack.pop();
            pair.forEach(PlatformTransactionManager::rollback);
        }
        THREAD_LOCAL.remove();
    }

}
